import torch
import numpy as np
import matplotlib.pyplot as plt

# y = 3*x**2 + 2 * x + 2

x_data = [1.0, 2.0, 3.0]
y_data = [7.0,18.0,35.0]

w1 = torch.tensor([1.0])
w1.requires_grad_(True)
b = torch.tensor([0.0])
b.requires_grad_(True)
w2 = torch.tensor([1.0])
w2.requires_grad_(True)

'''
    每次做的事情就是先计算出loss，然后求出backward即可，反向传播
'''

def forward(x):
#     return x * w + b
    return w1 * x ** 2 + w2 * x + b


def loss(x, y):
    y_pred = forward(x)
    return (y_pred - y) * (y_pred - y)

loss_list = []
epoch_list = []

print("Predict(before training)", 4, forward(4))

for epoch in range(1000):
    for x, y in zip(x_data, y_data):
        l = loss(x, y)
        l.backward()
        print("\tgrad:", x, y, w1.grad.item(), w2.grad.item(), b.grad.item())
        w1.data -= w1.grad.data * 0.01
        w2.data -= w2.grad.data * 0.01
        b.data -= b.grad.data * 0.01

        w1.grad.data.zero_()
        w2.grad.data.zero_()
        b.grad.data.zero_()

        epoch_list.append(epoch)
        loss_list.append(l.item())
    # print("progess:",epoch,l.item())


print(f"\n优化后的参数: w1={w1.item():.4f}, w2={w2.item():.4f}, b={b.item():.4f}")
print("Predict(after training):", 4, forward(4).item())  # 预期输出应为3*16 + 2*4 + 2 = 58

plt.plot(epoch_list, loss_list)
plt.xlabel("epoch")
plt.ylabel("loss")
plt.show()

